import pandas as pd
from sklearn.preprocessing import LabelEncoder
import os


def have_train_preprocess_data():
    return os.path.exists('../data/train_preprocess.csv')


def have_test_preprocess_data():
    return os.path.exists('../data/test_preprocess.csv')


def load_train_preprocess_data():
    return pd.read_csv('../data/train_preprocess.csv')


def load_test_preprocess_data():
    return pd.read_csv('../data/test_preprocess.csv')


def train_data_analysis_encode_save():
    data = pd.read_csv('../data/train.csv')
    data = preprocess_data_change(data)
    data.to_csv('../data/train_preprocess.csv', index=False)
    return data


def test_data_analysis_encode_save():
    data = pd.read_csv('../data/test2.csv')
    data = preprocess_data_change(data)
    data.to_csv('../data/test_preprocess.csv', index=False)
    return data


def preprocess_data_change(data):
    # 去掉over18和standardhours以及员工编码
    data = data.drop(['Over18', 'StandardHours', 'EmployeeNumber'], axis=1)
    data = map_business_change(data)
    data = map_department_change(data)
    data['EducationField'] = LabelEncoder().fit_transform(data['EducationField'])
    # 热编码处理性别和是否加班列
    data = pd.get_dummies(data, columns=['Gender', 'OverTime'], drop_first=True)
    data['JobRole'] = LabelEncoder().fit_transform(data['JobRole'])
    data = map_maritalStatus_change(data)
    # Attrition默认是在最后一列，但train数据集中是第一列，并且测试集数据集处理完成编码后，Attrition就不再是最后一列了，所以要移动到第一列
    # 将 'Attrition' 列提取出来
    cols = ['Attrition'] + [col for col in data.columns if col != 'Attrition']
    # 重新排序 DataFrame 的列
    data = data[cols]
    return data


def map_business_change(data):
    # 手动定义映射关系
    mapping = {'Non-Travel': 0, 'Travel_Rarely': 1, 'Travel_Frequently': 2}
    # 使用 map() 替换，默认
    data['BusinessTravel'] = data['BusinessTravel'].map(mapping)
    return data


def map_department_change(data):
    # 此处也一样采用映射方法
    mapping = {'Human Resources': 1, 'Research & Development': 2, 'Sales': 3}
    data['Department'] = data['Department'].map(mapping)
    return data


def map_maritalStatus_change(data):
    mapping = {'Divorced': 0, 'Single': 1, 'Married': 2}
    data['MaritalStatus'] = data['MaritalStatus'].map(mapping)
    return data
